import os
import argparse
import struct # Potentially useful for more complex byte operations later, but not strictly needed for just zeros.

# --- 观天道项目内存配置 (必须与Rust代码中的mem_manager.rs保持一致) ---
SYS_SEGMENT_SIZE = 1024  # 1KB for System Segment
PUB_SEGMENT_WIDTH = 128
PUB_SEGMENT_HEIGHT = 128
PUB_SEGMENT_SIZE = PUB_SEGMENT_WIDTH * PUB_SEGMENT_HEIGHT  # 128*128 = 16384 bytes
TOTAL_MEMORY_SIZE = SYS_SEGMENT_SIZE + PUB_SEGMENT_SIZE  # Total 1024 + 16384 = 17408 bytes
# --- 配置结束 ---

def initialize_mem_img(output_path, force_overwrite=False):
    """
    初始化 mem.img 文件，填充所有字节为 0。
    """
    if os.path.exists(output_path):
        if force_overwrite:
            print(f"WARNING: File '{output_path}' already exists. Overwriting as --force was used.")
            # os.remove(output_path) # Not strictly necessary, 'wb' mode truncates
        else:
            print(f"ERROR: File '{output_path}' already exists. Use -f or --force to overwrite.")
            return False

    print(f"Initializing mem.img at '{output_path}'...")
    print(f"  System Segment Size: {SYS_SEGMENT_SIZE} bytes")
    print(f"  Public Segment Size: {PUB_SEGMENT_WIDTH}x{PUB_SEGMENT_HEIGHT} pixels = {PUB_SEGMENT_SIZE} bytes")
    print(f"  Total Memory Size: {TOTAL_MEMORY_SIZE} bytes ({TOTAL_MEMORY_SIZE / 1024:.2f} KB)")

    # 创建一个指定大小的字节数组，所有元素默认为0
    initial_data = bytearray(TOTAL_MEMORY_SIZE)

    try:
        with open(output_path, 'wb') as f:
            f.write(initial_data)
        print(f"Successfully initialized '{output_path}'.")
        return True
    except IOError as e:
        print(f"ERROR: Could not write to '{output_path}': {e}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Initialize the mem.img file for the Guan Tian Dao project. "
                    "Ensures consistent memory layout with Rust backend."
    )
    parser.add_argument(
        "output_path",
        nargs="?",  # Makes the argument optional
        default="mem.img",
        help="Path to the mem.img file to be initialized (default: mem.img)",
    )
    parser.add_argument(
        "-f", "--force",
        action="store_true",
        help="Force overwrite if mem.img already exists.",
    )

    args = parser.parse_args()

    initialize_mem_img(args.output_path, args.force)